package com.zx.lc.rag;

import static java.util.stream.Collectors.joining;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.springframework.stereotype.Service;

@Service
public class _2_LowLevelNativeRagService {

    /**
     * 模拟 langchain4j 底层是如何实现模型与向量 配合实现rag的。
     * @param args
     */
    public static void main(String[] args) {
        //1. 初始化大语言模型，
        ChatLanguageModel chatModel = OllamaChatModel.builder()
                .baseUrl("http://192.168.3.99:11434")
                .timeout(Duration.ofSeconds(3600))
                .modelName("qwen2.5:7b")
                .build();

        // 2. 初始化嵌入模型：langchain4j支持的向量模型：https://docs.langchain4j.dev/category/embedding-models/
        OllamaEmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://192.168.3.99:11434")
                .modelName("bge-m3")
                .build();

        // 3.初始化向量数据库。使用内存模拟向量数据库
        // langchain4j支持的向量存储：https://docs.langchain4j.dev/integrations/embedding-stores/
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        //  读取数据文件。
        //  1. 支持多种来源的数据解析：文件系统, URL, Amazon S3, Azure Blob存储, GitHub, 腾讯云 COS等
        //  2. 支持多种文件的解析：txt、pdf、doc、xls、ppt
        List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("D:\\llm\\rag-demo");

        // 将文件进行拆分为更小的块，方面精准的进行查询，而不是每次豆包珍整个文章传给模型，降低传输的数据量
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        for (Document document : documents) {
            List<TextSegment> segments = splitter.split(document);
            // 使用嵌入模型，将拆分后的文件进行向量化
            Response<List<Embedding>> embeddings = embeddingModel.embedAll(segments);
            // 将向量化的数据存入向量数据库中
            embeddingStore.addAll(embeddings.content(), segments);
        }

        // 问题
        String question = "陆瑶 是谁？帮我查下这个人的信息。";
        // 将文件向量化
        Embedding questionEmbedding = embeddingModel.embed(question).content();
        // 构建向量的查询条件
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(questionEmbedding)
                .maxResults(3)
                .minScore(0.7)
                .build();

        // 向 向量数据库查询问题
        List<EmbeddingMatch<TextSegment>> relevantEmbeddings = embeddingStore.search(embeddingSearchRequest).matches();

        // 创建一个提示词，对结果进行格式化
        PromptTemplate promptTemplate = PromptTemplate.from(
                "Answer the following question to the best of your ability:\n"
                        + "\n"
                        + "Question:\n"
                        + "{{question}}\n"
                        + "\n"
                        + "Base your answer on the following information:\n"
                        + "{{information}}");
        // 格式化内容
        String information = relevantEmbeddings.stream()
                .map(match -> match.embedded().text())
                .collect(joining("\n\n"));

        // 构建向模型的请求参数
        Map<String, Object> variables = new HashMap<>();
        variables.put("question", question);
        variables.put("information", information);

        Prompt prompt = promptTemplate.apply(variables);

        ChatResponse response = chatModel.chat(prompt.toUserMessage());
        System.out.println("tokenUsage：" + response.tokenUsage());
        System.out.println("metadata：" + response.metadata());
        System.out.println("aiMessage：" + response.aiMessage().text());
    }


}
